package com.zyx.kafkademo.consumer;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.zyx.kafkademo.utils.TimeUtils;
import org.apache.kafka.clients.consumer.*;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.StringDeserializer;

import java.time.Duration;
import java.time.Instant;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author Yaxi.Zhang
 * @since 2021/8/18 10:36
 * desc: 指定发送时间消费信息
 */
public class SeekTimeDemo {
    public static final String BROKER_LIST = "bigdata022:9092,bigdata023:9092,bigdata024:9092";
    public static final String TOPIC = "HK_TRADE_TICKER";
    public static final String GROUP_ID = "test_consumer_210818";

    public static Properties initConfig() {
        Properties props = new Properties();
        props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG,
                StringDeserializer.class.getName());
        props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG,
                StringDeserializer.class.getName());
        props.put(ConsumerConfig.ENABLE_AUTO_COMMIT_CONFIG, "false");
        props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, BROKER_LIST);
        props.put(ConsumerConfig.GROUP_ID_CONFIG, GROUP_ID);
        return props;
    }

    public static void main(String[] args) {

        Properties props = initConfig();

        try (KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props)) {
            consumer.subscribe(Collections.singletonList(TOPIC));

            long start = Instant.now().toEpochMilli();
            Set<TopicPartition> assignment = new HashSet<>();

            while (assignment.size() == 0) {
                // 如果没有分配到分区,就一直循环下去
                consumer.poll(Duration.ofMillis(100));
                assignment = consumer.assignment();
            }

            long end = Instant.now().toEpochMilli();

            System.out.println("已获取主题分区信息,costing " + (end - start) + "ms...");
            System.out.println("主题分区信息:" + assignment);


            // 指定分区对应的时间戳
            Map<TopicPartition, Long> map = new HashMap<>();
            for (TopicPartition tp : assignment) {
                map.put(tp, 1628841598000L);
            }

            Map<TopicPartition, OffsetAndTimestamp> offsets = consumer.offsetsForTimes(map);

            // 通过时间戳获取分区的offset从而获取数据
            for (TopicPartition topicPartition : offsets.keySet()) {
                OffsetAndTimestamp offsetAndTimestamp = offsets.get(topicPartition);
                if (offsetAndTimestamp != null) {
                    System.out.println(topicPartition + "对应的偏移量为:" + offsetAndTimestamp.offset());
                    consumer.seek(topicPartition, offsetAndTimestamp.offset());
                }
            }

            AtomicInteger index = new AtomicInteger(1);

            while (true) {
                System.out.println("index:" + index.intValue());
                // 消费信息
                ConsumerRecords<String, String> records = consumer.poll(Duration.ofMillis(1000));
                for (ConsumerRecord<String, String> record : records) {
                    String value = record.value();
                    JSONObject jsonObject = JSONObject.parseObject(value);
                    /*Long ts = jsonObject.getLong("time");
                    String formatTimeStr = TimeUtils.getFormatTimeStr(ts);

                    JSONArray jsonArray = jsonObject.getJSONArray("extraData");

                    for (Object data : jsonArray) {
                        JSONObject ticker = JSONObject.parseObject(JSON.toJSONString(data));
                        String ai = ticker.getString("ai");
                        if ("06666.HK".equals(ai)) {
                            Double price = ticker.getDouble("p");
                            Long quant = ticker.getLong("q");
                            System.out.println("partition:" + record.partition() + ",offset:" + record.offset() + ",ai:" + ai + ",price:" + price + ",quant:" + quant + "," + "time:" + formatTimeStr);
                        }

                    }*/

                    Long ts = jsonObject.getLong("sendTime");
                    String formatTimeStr = TimeUtils.getFormatTimeStr(ts);

                    Integer secCode = jsonObject.getInteger("securityCode");

                    if (secCode == 6666) {
                        Double price = jsonObject.getDouble("price");
                        Long quant = jsonObject.getLong("aggregateQuantity");


                        System.out.println("partition:" + record.partition() + ",offset:" + record.offset()
                                + ",ai:" + secCode + ",price:" + price + ",quant:" + quant + "," + "time:" + formatTimeStr + ",sendtime:" + ts);
                    }


                    // 输出信息
                    // System.out.println(record.offset() + ":" + record.value() + " ==> " + formatTimeStr);
                }

                int cur = index.incrementAndGet();
                if (cur >= 200) {
                    break;
                }
            }
        } catch (Exception exception) {
            exception.printStackTrace();
        }

    }

}
